import os
import pandas as pd
import tensorflow as tf
import tensorflow.contrib.estimator as tce
from tensorflow.python.data import Dataset
import numpy as np
from sklearn import metrics
import math

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
data = pd.read_csv('csv.csv')


def my_input_func(features, targets, batch_size=1, shuffle=True, num_epochs=None):
    features = {key: np.array(value) for key, value in dict(features).items()}
    ds = Dataset.from_tensor_slices((features, targets))
    ds = ds.batch(batch_size).repeat(num_epochs)
    if shuffle:
        ds = ds.shuffle(buffer_size=10000)
    features, labels = ds.make_one_shot_iterator().get_next()
    return features, labels


def predict_median_house_value():
    '''
    步骤：
        1. 定义输入
        2. 定义目标
        3. 配置训练模型
        4.
    '''
    feature = data[['total_rooms']]
    feature_col = [tf.feature_column.numeric_column('total_rooms')]

    target = data['median_house_value'] / 1000.0

    optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.0000001)
    optimizer = tce.clip_gradients_by_norm(optimizer, 5.0)

    liner_predict = tf.estimator.LinearRegressor(
        feature_columns=feature_col,
        optimizer=optimizer
    )

    '''训练'''
    liner_predict.train(input_fn=lambda: my_input_func(feature, target), steps=100)
    '''预测'''
    prediction_input_func = lambda: my_input_func(feature, target, num_epochs=1, shuffle=False)
    predictions = liner_predict.predict(input_fn=prediction_input_func)
    predictions = np.array([item['predictions'][0] for item in predictions])

    '''评估模型'''
    mse = metrics.mean_squared_error(predictions, target)
    rmse = math.sqrt(mse)

    print('mse:{},rmse:{}'.format(mse, rmse))

    # for item in predictions:
    #     print('{},{},{}'.format(type(item),item,item['predictions'][0]))
    #     '''
    #     <class 'dict'>,{'predictions': array([0.84449947], dtype=float32)}
    #     <class 'dict'>,{'predictions': array([0.7844995], dtype=float32)}
    #     '''


def practice_features():
    feature = data[['total_rooms']]
    '''key = total_rooms,value = 数据内容'''
    f = {key: np.array(value) for key, value in dict(feature).items()}
    return f


def test():
    f = practice_features()
    assert isinstance(f, dict)
    assert 'total_rooms' in f
    assert f['total_rooms'][0] == 5612.0
    assert f['total_rooms'][0] == 5612.
    assert f['total_rooms'][0] == 5612


if __name__ == '__main__':
    predict_median_house_value()
    # test()
    print('*' * 50)
    print('test passed'.center(50, '*'))
    print('*' * 50)
